"""
请求限流器 - 使用令牌桶算法控制API调用频率
"""
import time
import threading
from typing import Optional


class TokenBucket:
    """令牌桶算法实现"""
    
    def __init__(self, rate: float, capacity: float):
        """
        初始化令牌桶
        
        Args:
            rate: 令牌生成速率（每秒）
            capacity: 桶容量（突发请求数）
        """
        self.rate = rate
        self.capacity = capacity
        self.tokens = capacity
        self.last_update = time.time()
        self.lock = threading.Lock()
    
    def consume(self, tokens: int = 1) -> bool:
        """
        消费令牌
        
        Args:
            tokens: 需要消费的令牌数
            
        Returns:
            是否成功消费令牌
        """
        with self.lock:
            now = time.time()
            # 计算新增令牌数
            elapsed = now - self.last_update
            self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
            self.last_update = now
            
            # 检查是否有足够令牌
            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            return False
    
    def wait_for_token(self, tokens: int = 1, timeout: Optional[float] = None) -> bool:
        """
        等待令牌可用
        
        Args:
            tokens: 需要的令牌数
            timeout: 超时时间（秒），None表示无限等待
            
        Returns:
            是否成功获取令牌
        """
        start_time = time.time()
        
        while True:
            if self.consume(tokens):
                return True
            
            # 检查超时
            if timeout is not None:
                elapsed = time.time() - start_time
                if elapsed >= timeout:
                    return False
            
            # 短暂休眠后重试
            time.sleep(0.01)


class RateLimiter:
    """API请求限流器"""
    
    def __init__(self, requests_per_second: float = 10, burst: float = 20):
        """
        初始化限流器
        
        Args:
            requests_per_second: 每秒允许的请求数
            burst: 突发请求容量
        """
        self.bucket = TokenBucket(rate=requests_per_second, capacity=burst)
        self.enabled = True
    
    def acquire(self, timeout: Optional[float] = None) -> bool:
        """
        获取请求许可
        
        Args:
            timeout: 超时时间（秒）
            
        Returns:
            是否成功获取许可
        """
        if not self.enabled:
            return True
        
        return self.bucket.wait_for_token(tokens=1, timeout=timeout)
    
    def disable(self):
        """禁用限流"""
        self.enabled = False
    
    def enable(self):
        """启用限流"""
        self.enabled = True
